import os
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

# 本地模型文件夹路径（替换为你的实际路径）
local_model_path = "/opt/model/roberta-base-finetuned-dianping-chinese"

def check_model_files(path):
    """修正文件检查逻辑，兼容pytorch_model.bin"""
    required_files = [
        "config.json",           # 模型配置文件（必须）
        "tokenizer_config.json", # 分词器配置（必须）
        "vocab.txt"              # 分词器词表（必须）
    ]
    # 权重文件可能是以下两种格式之一（满足其一即可）
    weight_files = [
        "pytorch_model.bin",     # PyTorch常用格式
        "model.safetensors"      # 另一种安全格式
    ]

    # 检查路径是否存在
    if not os.path.exists(path):
        return False, f"路径不存在: {path}"

    # 检查必要文件
    missing_required = [f for f in required_files if not os.path.exists(os.path.join(path, f))]
    if missing_required:
        return False, f"缺少必要配置文件: {', '.join(missing_required)}"

    # 检查权重文件（至少存在一个）
    weight_exists = any(os.path.exists(os.path.join(path, f)) for f in weight_files)
    if not weight_exists:
        return False, f"缺少模型权重文件，需至少存在一个: {', '.join(weight_files)}"

    return True, "模型文件完整"

try:
    # 验证模型文件
    is_valid, message = check_model_files(local_model_path)
    if not is_valid:
        raise ValueError(message)

    # 加载模型和分词器
    tokenizer = AutoTokenizer.from_pretrained(
        local_model_path,
        trust_remote_code=True,
        local_files_only=True
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        local_model_path,
        trust_remote_code=True,
        local_files_only=True
    )

    # 测试情感分析
    classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
    result = classifier("今儿上海可真冷啊")
    print("情感分析结果：", result)

except Exception as e:
    print(f"错误：{e}")